set.seed(1782318297)

library(tidyverse)

# Robstuness tests 2

## baseline confounders
pre_cov_robust <- c("age", "female", "education", "income", "race", "urban",
             "country")

##Mediator and outcome equations

m_form_robust <- sat ~ age + female + education + income + race + urban + country + vic + aoj11 +
  b13 + b18 + b21 + b21a + b31 + l1 + wave

y_form_robust <- sup ~ sat + age + female + education + income + race + urban + country + vic + 
  aoj11*sat + b13*sat + b18*sat + b21*sat + b21a*sat + b31*sat + l1*sat + wave


## Models for the post-treatment confounder 

m1_robust <- lm(aoj11 ~ age + female + education + income + race + urban + country + vic, weights = weights,
         data = nn_data)

m2_robust <- lm(b13 ~ age + female + education + income + race + urban + country + vic, weights = weights,
         data = nn_data)

m3_robust <- lm(b18 ~ age + female + education + income + race + urban + country + vic, weights = weights,
         data = nn_data)

m4_robust <- lm(b21 ~ age + female + education + income + race + urban + country + vic, weights = weights,
         data = nn_data)

m5_robust <- lm(b21a ~ age + female + education + income + race + urban + country + vic, weights = weights,
         data = nn_data)

m6_robust <- lm(b31 ~ age + female + education + income + race + urban + country + vic, weights = weights,
         data = nn_data)

m7_robust <- lm(l1 ~ age + female + education + income + race + urban + country + vic, weights = weights,
         data = nn_data)

m8_robust <- lm(aoj11 ~ age + female + education + income + race + urban + country + vic, weights = weights,
         data = cem_data)

m9_robust <- lm(b13 ~ age + female + education + income + race + urban + country + vic, weights = weights,
         data = cem_data)

m10_robust <- lm(b18 ~ age + female + education + income + race + urban + country + vic, weights = weights,
          data = cem_data)

m11_robust <- lm(b21 ~ age + female + education + income + race + urban + country + vic, weights = weights,
          data = cem_data)

m12_robust <- lm(b21a ~ age + female + education + income + race + urban + country + vic, weights = weights,
          data = cem_data)

m13_robust <- lm(b31 ~ age + female + education + income + race + urban + country + vic, weights = weights,
          data = cem_data)

m14_robust <- lm(l1 ~ age + female + education + income + race + urban + country + vic, weights = weights,
          data = cem_data)

m15_robust <- lm(aoj11 ~ age + female + education + income + race + urban + country + vic, weights = weights,
          data = full_data)

m16_robust <- lm(b13 ~ age + female + education + income + race + urban + country + vic, weights = weights,
          data = full_data)

m17_robust <- lm(b18 ~ age + female + education + income + race + urban + country + vic, weights = weights,
          data = full_data)

m18_robust <- lm(b21 ~ age + female + education + income + race + urban + country + vic, weights = weights,
          data = full_data)

m19_robust <- lm(b21a ~ age + female + education + income + race + urban + country + vic, weights = weights,
          data = full_data)

m20_robust <- lm(b31 ~ age + female + education + income + race + urban + country + vic, weights = weights,
          data = full_data)

m21_robust <- lm(l1 ~ age + female + education + income + race + urban + country + vic, weights = weights,
          data = full_data)



## RWR estimation

install.packages("remotes")
remotes::install_github("xiangzhou09/rwrmed")

library(rwrmed)

fit_nn_robust <- rwrmed(treatment = "vic", pre_cov = pre_cov_robust, zmodels = list(m1_robust, m2_robust, m3_robust, m4_robust, m5_robust, m6_robust, m7_robust),
                 y_form = y_form_robust, m_form = m_form_robust, weights = weights, data = nn_data)

fit_full_robust <- rwrmed(treatment = "vic", pre_cov = pre_cov_robust, zmodels = list(m15_robust, m16_robust, m17_robust, m18_robust, m19_robust, m20_robust, m21_robust),
                   y_form = y_form_robust, m_form = m_form_robust, weights = weights, data = full_data)

fit_cem_robust <- rwrmed(treatment = "vic", pre_cov = pre_cov_robust, zmodels = list(m8_robust, m9_robust, m10_robust, m11_robust, m12_robust, m13_robust, m14_robust),
                  y_form = y_form_robust, m_form = m_form_robust, weights = weights, data = cem_data)

##Effect decomposition


###Nearest Neighbor


out_nn_robust <- decomp(fit_nn_robust, rep = 500)

print(out_nn_robust, digits = 2)

twocomp_nn_robust <- out_nn_robust[["twocomp"]]

twocomp_nn_robust <- as.data.frame(twocomp_nn_robust)

twocomp_nn_robust <- twocomp_nn_robust %>% 
  rownames_to_column(var = "effect_type") %>% 
  mutate(match = "Nearest Neighbor")

###Full Matching

out_full_robust <- decomp(fit_full_robust, rep = 500)

print(out_full_robust, digits = 2)

twocomp_full_robust <- out_full_robust[["twocomp"]]

twocomp_full_robust <- as.data.frame(twocomp_full_robust)

twocomp_full_robust <- twocomp_full_robust %>% 
  rownames_to_column(var = "effect_type") %>% 
  mutate(match = "Full Matching")


###CEM

out_cem_robust <- decomp(fit_cem_robust, rep = 500)
print(out_cem_robust, digits = 2)

twocomp_cem_robust <- out_cem_robust[["twocomp"]]

twocomp_cem_robust <- as.data.frame(twocomp_cem_robust)

twocomp_cem_robust <- twocomp_cem_robust %>% 
  rownames_to_column(var = "effect_type") %>% 
  mutate(match = "CEM")

###Figure 5


twocomp_robust <- rbind(twocomp_nn_robust, twocomp_full_robust, twocomp_cem_robust)

library(showtext)

font_add("Cambria", "cambria.ttc")

font_families()

showtext_auto()

twocomp_plot_robust <- twocomp_robust %>% 
  ggplot(aes(x = Estimate, y = effect_type)) +
  geom_point() +
  geom_pointrange(aes(xmin = `2.5% Perc`, xmax = `97.5% Perc`)) +
  geom_vline(xintercept = 0, linetype = "dashed", color = "black") +
  labs(y = "", x = "") +
  theme_minimal() +
  theme(text=element_text(size=18,  family= "Cambria")) +
  facet_grid(~match)

###Support for coup (Footnote 1)

y_form_c <- coup ~ sat + age + female + education + income + race + urban + country + vic + 
  aoj11*sat + b13*sat + b18*sat + b21*sat + b21a*sat + b31*sat + l1*sat + wave


fit_full_c <- rwrmed(treatment = "vic", pre_cov = pre_cov_robust, zmodels = list(m15_robust, m16_robust, m17_robust, m18_robust,
                                                                                        m19_robust, m20_robust, m21_robust),
                   y_form = y_form_c, m_form = m_form_robust, weights = weights, data = full_data)

out_full_c <- decomp(fit_full_c, rep = 500)


twocomp_full_c <- out_full_c[["twocomp"]]

write.csv2(twocomp_full_c, "supcoup.csv")
